import os
import uuid
from bs4 import BeautifulSoup
from collections import deque
from logger import get_logger
from graphviz import Digraph
from torch_geometric.data import Data
import torch
logger = get_logger()


class H2G:
    """
        将html转换为图数据
    """

    def __init__(self, html_path, is_debug: bool = False) -> None:
        """
        初始化

        :param html_path : html文件路经
        :param is_debug : debug
        :returns:
        """
        self.html_path = html_path
        # ===debug===
        is_save = True if is_debug else False
        is_render = True if is_debug else False

        # ===========
        self._annotated_soup = self._annotate_node_id(is_save=is_save, is_render=is_render)  # 每个节点有唯一标识id
        self._graph = self._html_to_graph(is_render=is_render)  # 转换为图数据

    def get_graph(self) -> Data:
        return self._graph

    def _annotate_node_id(self, is_save: bool = False, is_render: bool = True) -> BeautifulSoup:
        """
        给层次遍历HTML树并给每个节点标注唯一id

        :param is_save : 是否保存标注好的html文件
        :param is_render : 是否渲染html树
        :returns: 
        """
        # 读取 HTML 文件
        with open(self.html_path, "r") as f:
            html = f.read()

        # 解析 HTML
        soup = BeautifulSoup(html, 'html.parser')

        id = 0
        queue = deque([soup.html])
        while queue:
            id += 1
            node = queue.popleft()
            node['ann_id'] = id
            for child in node.children:
                if child.name is not None:  # 只将标签节点加入队列
                    queue.append(child)
        if is_save:
            # 生成一个随机的 UUID
            unique_id = str(uuid.uuid4())[:5]
            save_path = f'./ann_id_html/{unique_id}.html'
            if not os.path.exists("./ann_id_html"):
                os.makedirs("./ann_id_html")
            try:
                with open(save_path, "w") as f:
                    f.write(soup.prettify())
                logger.info(f"文件保存在 {save_path}")
            except Exception as e:
                logger.error(str(e))
        if is_render:
            dot = Digraph()
            queue = deque([soup.html])
            while queue:
                node = queue.popleft()
                dot.node(str(node.name) + "-" + str(node['ann_id']))
                for child in node.children:
                    if child.name is not None:  # 只将标签节点加入队列
                        queue.append(child)
                        dot.edge(str(node.name) + "-" + str(node['ann_id']), str(child.name) + "-" + str(child['ann_id']))
            self._render_html(dot)
        return soup

    def _html_to_graph(self, is_render: bool = False) -> Data:
        # x = []  # 节点特征 [num_nodes, num_node_features]
        edge_index = []  # 边的索引 [num_edges, 2]
        queue = deque([self._annotated_soup.html])
        num_nodes = 0
        while queue:
            num_nodes += 1
            node = queue.popleft()
            for child in node.children:
                if child.name is not None:  # 只将标签节点加入队列
                    queue.append(child)
                    edge_index.append([int(node['ann_id'])-1, int(child['ann_id'])-1])
        # print(edge_index)

        data = Data(x=torch.randn(num_nodes, 1), edge_index=torch.tensor(edge_index, dtype=torch.long))
        if is_render:
            dot = Digraph()
            for i in range(data.num_nodes):
                dot.node(str(i+1))
            for i in data.edge_index.tolist():
                dot.edge(str(i[0]+1), str(i[1]+1))
            self._render_graph(dot)
        return data

    def _render_html(self, dot: Digraph):
        # 生成一个随机的 UUID
        unique_id = str(uuid.uuid4())[:5]
        if not os.path.exists("./render_html"):
            os.makedirs("./render_html")
        render_file_path = f'./render_html/{unique_id}'
        try:
            dot.render(render_file_path, format="pdf")
            logger.info(f"渲染html文件已保存到 {render_file_path}.pdf")
        except Exception as e:
            logger.error(str(e))

    def _render_graph(self, dot: Digraph):
        # 生成一个随机的 UUID
        unique_id = str(uuid.uuid4())[:5]
        if not os.path.exists("./render_graph"):
            os.makedirs("./render_graph")
        render_file_path = f'./render_graph/{unique_id}'
        try:
            dot.render(render_file_path, format="pdf")
            logger.info(f"渲染graph文件已保存到 {render_file_path}.pdf")
        except Exception as e:
            logger.error(str(e))


if __name__ == '__main__':
    h2g = H2G("./demo.html")
    print(h2g.get_graph())
